// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/expression/decoded_args.h>
#include <pollux/expression/string_writer.h>
#include <pollux/expression/vector_function.h>
#include <pollux/functions/lib/utf8_utils.h>
#include <pollux/common/strings/string_impl.h>

namespace kumo::pollux::functions {
    namespace {
        class FromUtf8Function : public exec::VectorFunction {
        public:
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr & /* outputType */,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                const auto &input = args[0];

                // Optimize common cases:
                // (1) all-ASCII input;
                // (2) valid UTF-8 input;
                // (3) constant replacement character.

                exec::DecodedArgs decodedArgs(rows, args, context);
                auto &decodedInput = *decodedArgs.at(0);

                // Read the constant replacement if it exisits and verify that it is valid.
                bool constantReplacement =
                        args.size() == 1 || decodedArgs.at(1)->isConstantMapping();
                std::string constantReplacementValue = kReplacementChar;

                if (constantReplacement) {
                    if (args.size() > 1) {
                        auto &decodedReplacement = *decodedArgs.at(1);
                        try {
                            constantReplacementValue = getReplacementCharacter(
                                args[1]->type(), decodedReplacement, rows.begin());
                        } catch (const std::exception &) {
                            context.setErrors(rows, std::current_exception());
                            return;
                        }
                    }
                }

                // We can only do valid UTF-8 input optimization if replacement is valid and
                // constant otherwise we have to check replacement for each row.
                if (constantReplacement) {
                    if (input->loaded_vector()
                        ->as<SimpleVector<StringView> >()
                        ->computeAndSetIsAscii(rows)) {
                        // Input strings are all-ASCII.
                        toVarcharNoCopy(input, decodedInput, rows, context, result);
                        return;
                    }
                }

                auto firstInvalidRow = findFirstInvalidRow(decodedInput, rows);

                // We can only do this optimization if replacement is valid and
                // constant otherwise we have to check replacement for each row.
                if (constantReplacement) {
                    if (!firstInvalidRow.has_value()) {
                        // All inputs are valid UTF-8 strings.
                        toVarcharNoCopy(input, decodedInput, rows, context, result);
                        return;
                    }
                }

                BaseVector::ensure_writable(rows, VARCHAR(), context.pool(), result);
                auto flatResult = result->as<FlatVector<StringView> >();

                // Reserve string buffer capacity.
                size_t totalInputSize = 0;
                rows.applyToSelected([&](auto row) {
                    totalInputSize += decodedInput.value_at<StringView>(row).size();
                });

                flatResult->getBufferWithSpace(totalInputSize);

                if (constantReplacement) {
                    rows.applyToSelected([&](auto row) {
                        exec::StringWriter writer(flatResult, row);
                        auto value = decodedInput.value_at<StringView>(row);
                        if (row < firstInvalidRow) {
                            writer.append(value);
                            writer.finalize();
                        } else {
                            fixInvalidUtf8(value, constantReplacementValue, writer);
                        }
                    });
                } else {
                    auto &decodedReplacement = *decodedArgs.at(1);
                    context.applyToSelectedNoThrow(rows, [&](auto row) {
                        auto replacement =
                                getReplacementCharacter(args[1]->type(), decodedReplacement, row);
                        exec::StringWriter writer(flatResult, row);
                        auto value = decodedInput.value_at<StringView>(row);
                        if (row < firstInvalidRow) {
                            writer.append(value);
                            writer.finalize();
                        } else {
                            fixInvalidUtf8(value, replacement, writer);
                        }
                    });
                }
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                return {
                    // varbinary -> varchar
                    exec::FunctionSignatureBuilder()
                    .returnType("varchar")
                    .argumentType("varbinary")
                    .build(),

                    // varbinary, bigint -> varchar
                    exec::FunctionSignatureBuilder()
                    .returnType("varchar")
                    .argumentType("varbinary")
                    .argumentType("bigint")
                    .build(),

                    // varbinary, varchar -> varchar
                    exec::FunctionSignatureBuilder()
                    .returnType("varchar")
                    .argumentType("varbinary")
                    .argumentType("varchar")
                    .build(),
                };
            }

        private:
            static const std::string kReplacementChar;

            static std::string codePointToString(int64_t codePoint) {
                std::string result;
                result.resize(4);
                stringImpl::codePointToString(result, codePoint);
                return result;
            }

            std::string getReplacementCharacter(
                const TypePtr &type,
                DecodedVector &decoded,
                vector_size_t row) const {
                if (type->is_bigint()) {
                    return codePointToString(decoded.value_at<int64_t>(row));
                }

                auto replacement = decoded.value_at<StringView>(row);
                if (!replacement.empty()) {
                    int32_t codePoint;
                    auto charLength = tryGetUtf8CharLength(
                        replacement.data(), replacement.size(), codePoint);
                    POLLUX_USER_CHECK_GT(
                        charLength, 0, "Replacement is not a valid UTF-8 character");
                    POLLUX_USER_CHECK_EQ(
                        charLength,
                        replacement.size(),
                        "Replacement string must be empty or a single character");
                }
                return replacement;
            }

            /// Returns first row that contains invalid UTF-8 string or std::nullopt if
            /// all rows are valid.
            static std::optional<vector_size_t> findFirstInvalidRow(
                const DecodedVector &decodedInput,
                const SelectivityVector &rows) {
                std::optional<vector_size_t> firstInvalidRow;
                rows.testSelected([&](auto row) {
                    auto value = decodedInput.value_at<StringView>(row);

                    int32_t pos = 0;
                    while (pos < value.size()) {
                        int32_t codePoint;
                        auto charLength = tryGetUtf8CharLength(
                            value.data() + pos, value.size() - pos, codePoint);
                        if (charLength < 0) {
                            firstInvalidRow = row;
                            return false;
                        }

                        pos += charLength;
                    }

                    return true;
                });
                return firstInvalidRow;
            }

            void toVarcharNoCopy(
                const VectorPtr &input,
                DecodedVector &decodedInput,
                const SelectivityVector &rows,
                const exec::EvalCtx &context,
                VectorPtr &result) const {
                VectorPtr localResult;
                if (decodedInput.isConstantMapping()) {
                    auto value = decodedInput.value_at<StringView>(rows.begin());
                    localResult = std::make_shared<ConstantVector<StringView> >(
                        context.pool(), rows.end(), false, VARCHAR(), std::move(value));
                } else if (decodedInput.isIdentityMapping()) {
                    auto flatInput = decodedInput.base()->as_flat_vector<StringView>();

                    auto stringBuffers = flatInput->stringBuffers();
                    POLLUX_CHECK_LE(rows.end(), flatInput->size());
                    localResult = std::make_shared<FlatVector<StringView> >(
                        context.pool(),
                        VARCHAR(),
                        nullptr,
                        rows.end(),
                        flatInput->values(),
                        std::move(stringBuffers));
                } else {
                    auto base = decodedInput.base();
                    if (base->is_constant_encoding()) {
                        auto value = decodedInput.value_at<StringView>(rows.begin());
                        localResult = std::make_shared<ConstantVector<StringView> >(
                            context.pool(), rows.end(), false, VARCHAR(), std::move(value));
                    } else {
                        auto flatBase = base->as_flat_vector<StringView>();
                        auto stringBuffers = flatBase->stringBuffers();

                        auto values = AlignedBuffer::allocate<StringView>(
                            rows.end(), context.pool(), StringView());
                        auto *rawValues = values->asMutable<StringView>();
                        rows.applyToSelected([&](auto row) {
                            rawValues[row] = decodedInput.value_at<StringView>(row);
                        });

                        localResult = std::make_shared<FlatVector<StringView> >(
                            context.pool(),
                            VARCHAR(),
                            nullptr,
                            rows.end(),
                            std::move(values),
                            std::move(stringBuffers));
                    }
                }

                context.moveOrCopyResult(localResult, rows, result);
            }

            void fixInvalidUtf8(
                StringView input,
                const std::string &replacement,
                exec::StringWriter &fixedWriter) const {
                if (input.empty()) {
                    fixedWriter.setEmpty();
                    return;
                }

                int32_t pos = 0;
                while (pos < input.size()) {
                    int32_t codePoint;
                    auto charLength = tryGetUtf8CharLength(
                        input.data() + pos, input.size() - pos, codePoint);
                    if (charLength > 0) {
                        fixedWriter.append(std::string_view(input.data() + pos, charLength));
                        pos += charLength;
                        continue;
                    }

                    if (!replacement.empty()) {
                        fixedWriter.append(replacement);
                    }

                    pos += -charLength;
                }

                fixedWriter.finalize();
            }
        };

        // static
        const std::string FromUtf8Function::kReplacementChar =
                codePointToString(0xFFFD);
    } // namespace

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_from_utf8,
        FromUtf8Function::signatures(),
        std::make_unique<FromUtf8Function>());
} // namespace kumo::pollux::functions
